package com.zfei.gfb.core;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.Callable;

/**
 * @author fgh
 */
public class Context {

    private static final Logger logger = LoggerFactory.getLogger(Context.class);

    public static final Context ROOT = new Context(null);

    private static final ThreadLocal<Context> LOCAL = ThreadLocal.withInitial(() -> ROOT);

    public static Context current() {
        Context current = LOCAL.get();
        return (current == null ? ROOT : current);
    }

    public static Context getContext() {
        return LOCAL.get();
    }

    public static void removeContext() {
        LOCAL.remove();
    }

    private final Context parent;

    private final Map<String, Object> attachments;

    private Context(Context parent) {
        this.parent = parent;
        this.attachments = new LinkedHashMap<>();
    }

    private Context(Context parent, Map<String, Object> attachments) {
        this.parent = parent;
        this.attachments = attachments;
    }

    public Context fork() {
        return new Context(this);
    }

    /**
     * 多线程切换上下文
     *
     * @return 前一个上下文
     */
    public Context attach() {
        Context previous = current();
        LOCAL.set(this);
        return previous;
    }

    /**
     * 多线程还原上下文
     *
     * @param toAttach 将要还原的上下文
     */
    public void detach(Context toAttach) {
        if (toAttach == null) {
            logger.error("Context is null when detaching");
        } else if (toAttach.attach() != this) {
            logger.error("Context is not attached when detaching");
        }
    }

    /**
     * 多线程场景下判断是否当前上下文
     *
     * @return 结果
     */
    public boolean isCurrent() {
        return current() == this;
    }

    public Context getParent() {
        return parent;
    }

    public Object getAttachment(String key) {
        return attachments.get(key);
    }

    public Context setAttachment(String key, Object value) {
        if (value == null) {
            attachments.remove(key);
        } else {
            attachments.put(key, value);
        }
        return this;
    }

    public Context removeAttachment(String key) {
        attachments.remove(key);
        return this;
    }

    public Map<String, Object> getAttachments() {
        return attachments;
    }

    public Context setAttachments(Map<String, Object> attachment) {
        this.attachments.clear();
        if (attachment != null && attachment.size() > 0) {
            this.attachments.putAll(attachment);
        }
        return this;
    }

    public void clearAttachments() {
        this.attachments.clear();
    }

    public Object getGlobalAttachment(String key) {
        if (this != ROOT) {
            Object value;
            if (this.getParent() == ROOT) {
                value = this.getAttachment(key);
            } else {
                value = this.getParent().getGlobalAttachment(key);
            }
            return value;
        }
        return null;
    }

    public Context setGlobalAttachment(String key, Object value) {
        if (this != ROOT) {
            if (this.getParent() == ROOT) {
                this.setAttachment(key, value);
            } else {
                this.getParent().setGlobalAttachment(key, value);
            }
        }
        return this;
    }

    /**
     * Immediately run a {@link Runnable} with this context as the {@link #current} context.
     *
     * @param r {@link Runnable} to run.
     */
    public void run(Runnable r) {
        Context previous = attach();
        try {
            r.run();
        } finally {
            detach(previous);
        }
    }

    /**
     * Immediately call a {@link Callable} with this context as the {@link #current} context.
     *
     * @param c {@link Callable} to call.
     * @return result of call.
     */
    public <V> V call(Callable<V> c) throws Exception {
        Context previous = attach();
        try {
            return c.call();
        } finally {
            detach(previous);
        }
    }

    /**
     * Wrap a {@link Runnable} so that it executes with this context as the {@link #current} context.
     */
    public Runnable wrap(final Runnable r) {
        return () -> {
            Context previous = attach();
            try {
                r.run();
            } finally {
                detach(previous);
            }
        };
    }

    /**
     * Wrap a {@link Callable} so that it executes with this context as the {@link #current} context.
     */
    public <C> Callable<C> wrap(final Callable<C> c) {
        return () -> {
            Context previous = attach();
            try {
                return c.call();
            } finally {
                detach(previous);
            }
        };
    }
}
